package jamirr.apt;

import com.squareup.javapoet.*;
import jamirr.*;

import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.*;
import javax.lang.model.type.ArrayType;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeMirror;
import java.lang.annotation.Annotation;
import java.util.*;

public class MirrorElementProcessor extends AbstractElementProcessor {

    private static final String SUFFIX = "$Ref";

    @Override
    protected void processAnnotation(DeclaredType declaredType, ProcessorContext ctx) {

        var builder = TypeSpec.classBuilder("%s%s".formatted(ctx.getSimpleName(declaredType), SUFFIX))
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .superclass(ctx.getClassName(AnnotationReferenceImpl.class))
                .addAnnotation(AnnotationSpec.builder(Generated.class)
                        .addMember("value", "$S", "Auto generated by jamirr-apt. DO NOT EDIT!")
                        .build());

        // attributes
        var enclosedMethods = ctx.methodsIn(declaredType);
        // static
        var statik = CodeBlock.builder();
        boolean hasStatic = false;
        for (var e : enclosedMethods) {
            String name = e.getSimpleName().toString();

            var returnType = e.getReturnType();
            TypeName typeName = ClassName.get(returnType);
            // 静态字段
            builder.addField(FieldSpec.builder(ParameterizedTypeName.get(
                            ClassName.get("jamirr", "AnnotationAttributeBuilder"),
                            typeName.isPrimitive() ? typeName.box() : typeName
                    ), name)
                    .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                    .initializer("new $T($S)", ctx.getClassName(AnnotationAttributeBuilder.class), name)
                    .build());

            AnnotationValue defaultValue = e.getDefaultValue();

            if (defaultValue != null) {
                if (returnType instanceof ArrayType arrayType) {
                    // 静态方法
                    statik.addStatement("$L.setDefaultValue(new $T $L)", name, ClassName.get(arrayType), defaultValue);
                } else {
                    statik.addStatement("$L.setDefaultValue($L)", name, defaultValue);
                }
                hasStatic = true;
            }

            // @AliasFor
            AnnotationMirror aliasForAnnotation = null;
            List<? extends AnnotationMirror> attrAnnotations = e.getAnnotationMirrors();
            for (AnnotationMirror attrAnnotation : attrAnnotations) {
                if (ctx.isSame(attrAnnotation.getAnnotationType(), AliasFor.class.getCanonicalName())) {
                    aliasForAnnotation = attrAnnotation;
                    break;
                }
            }

            if (aliasForAnnotation != null) {
                var elementValues = aliasForAnnotation.getElementValues();
                elementValues.forEach((k, v) -> {
                    if (k.getSimpleName().toString().equals("attribute")) {
                        Object vv = v.getValue();
                        if (vv instanceof String aliasFor && !aliasFor.isEmpty()) {
                            statik.addStatement("$L.setAliasFor($S)", name, aliasFor);
                        }
                    } else if (k.getSimpleName().toString().equals("annotation")) {
                        Object vv = v.getValue();
                        if (vv instanceof TypeMirror tm && !ctx.isSame(tm, Annotation.class.getCanonicalName())) {
                            statik.addStatement("$L.setAliasForAnnotation($S)", name, ctx.getClassName(tm));
                        }
                    }
                });
            }
        }

        if (hasStatic) {
            builder.addStaticBlock(statik.build());
        }

        // getAnnotationName
        builder.addMethod(MethodSpec.methodBuilder("getAnnotationName")
                .addAnnotation(ctx.getClassName(Override.class))
                .addModifiers(Modifier.PUBLIC)
                .returns(ctx.getClassName(String.class))
                .addStatement("return $S", ctx.getFullyQualifiedName(declaredType))
                .build());

        // build
        MethodSpec.Builder build = MethodSpec.methodBuilder("build")
                .addAnnotation(ctx.getClassName(Override.class))
                .addModifiers(Modifier.PROTECTED)
                .addParameter(ClassName.get("jamirr", "AnnotationMirrorBuilder"), "builder")
                .returns(TypeName.VOID);

        for (var e : enclosedMethods) {
            String name = e.getSimpleName().toString();
            build.addStatement("builder.declareAttribute($L)", name);
        }

        var annotationMirrors = declaredType.asElement().getAnnotationMirrors();
        for (var annotationMirror : annotationMirrors) {
            String name = annotationMirror.getAnnotationType().toString();
            if (name.equals("jamirr.Mirror")
                    || name.startsWith("java.lang.annotation.")) {
                continue;
            }

            StringBuilder sb = new StringBuilder("builder.declaredAnnotation($S");
            List<Object> args = new ArrayList<>();
            args.add(name);
            for (var e : annotationMirror.getElementValues().entrySet()) {
                ExecutableElement key = e.getKey();
                AnnotationValue value = e.getValue();
                sb.append(", $S");
                args.add(key.getSimpleName().toString());
                sb.append(", $L");
                args.add(value);
            }
            sb.append(")");

            build.addStatement(sb.toString(), args.toArray());
        }
        builder.addMethod(build.build());

        TypeSpec typeSpec = builder.build();

        // 生成类文件
        ctx.writeJavaFile(ctx.getPackageNameOf(declaredType), typeSpec);

        // 生成配置文件
        ctx.writeMetaInf(AnnotationReference.class, "%s$Ref".formatted(ctx.getFullyQualifiedName(declaredType)));
    }

    @Override
    protected void processElement(DeclaredType declaredType, ProcessorContext ctx) {

        TypeSpec.Builder builder = TypeSpec.classBuilder(ctx.getSimpleName(declaredType) + SUFFIX)
                .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                .superclass(ParameterizedTypeName.get(
                        ctx.getClassName(AbstractBeanMirrorAndReference.class),
                        ctx.getClassName(declaredType)))
                .addAnnotation(AnnotationSpec.builder(Generated.class)
                        .addMember("value", "$S", "Auto generated by jamirr-apt. DO NOT EDIT!")
                        .build());

        // $DECLARED_ANNOTATIONS
        BuildFormat declared = resolveDeclaredAnnotations(declaredType, ctx);
        builder.addField(FieldSpec.builder(
                        ArrayTypeName.of(ctx.getClassName(jamirr.AnnotationMirror.class)),
                        "$DECLARED_ANNOTATIONS")
                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                .initializer(declared.format(), declared.args())
                .build());

        // 解析类的属性、构造函数、方法等
        List<DeclaredConstructor> constructors = new ArrayList<>();
        int constructorIndex = 0;
        List<Element> properties = new ArrayList<>();
        Map<String, DeclaredMethod> methods = new LinkedHashMap<>();
        int methodIndex = 0;
        List<Element> enclosedElements = resolveEnclosedElements(declaredType);
        for (Element enclosedElement : enclosedElements) {
            if (enclosedElement.getModifiers().contains(Modifier.STATIC)) {
                continue;
            }
            if (ElementKind.FIELD == enclosedElement.getKind()) {
                properties.add(enclosedElement);
            } else if (ElementKind.CONSTRUCTOR == enclosedElement.getKind()) {
                if (enclosedElement.getModifiers().contains(Modifier.PUBLIC)) {
                    if (enclosedElement instanceof ExecutableElement constructor) {
                        if (constructor.toString().startsWith(ctx.getSimpleName(declaredType))) {
                            constructors.add(new DeclaredConstructor(constructorIndex++, constructor));
                        }
                    }
                }
            } else if (ElementKind.METHOD == enclosedElement.getKind()) {
                if (enclosedElement.getModifiers().contains(Modifier.PUBLIC)) {
                    if (enclosedElement instanceof ExecutableElement method) {
                        String name = method.getSimpleName().toString();
                        String returnType = method.getReturnType().toString();
                        int args = method.getParameters().size();
                        String argType = null;
                        if (args > 0) {
                            argType = method.getParameters().getFirst().asType().toString();
                        }
                        methods.put(
                                name,
                                new DeclaredMethod(methodIndex++, returnType, args, argType, method)
                        );
                    }
                }
            }
        }

        // $CONSTRUCTORS
        BuildFormat constructorFormat = resolveConstructorReferences(constructors, ctx);
        builder.addField(FieldSpec.builder(ArrayTypeName.of(ClassName.get("", "ConstructorRef")), "$CONSTRUCTOR_REFERENCES")
                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class)
                        .addMember("value", "$S", "unchecked")
                        .build())
                .initializer(constructorFormat.format(), constructorFormat.args())
                .build());

        // $PROPERTY_REFERENCES
        BuildFormat propertyFormat = resolvePropertyReferences(properties, methods, ctx);
        builder.addField(FieldSpec.builder(ArrayTypeName.of(ClassName.get("", "PropertyRef")), "$PROPERTY_REFERENCES")
                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class)
                        .addMember("value", "$S", "unchecked")
                        .build())
                .initializer(propertyFormat.format(), propertyFormat.args())
                .build());


        // $METHOD_REFERENCES
        BuildFormat methodFormat = resolveMethodReferences(methods.values(), ctx);
        builder.addField(FieldSpec.builder(ArrayTypeName.of(ClassName.get("", "MethodRef")), "$METHOD_REFERENCES")
                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class)
                        .addMember("value", "$S", "unchecked")
                        .build())
                .initializer(methodFormat.format(), methodFormat.args())
                .build());

        // 构造器
        builder.addMethod(MethodSpec.constructorBuilder()
                .addModifiers(Modifier.PUBLIC)
                .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class)
                        .addMember("value", "$S", "unchecked")
                        .build())
                .addStatement("super($T.class, $L, $L, $L, $L)", ClassName.get(declaredType), "$DECLARED_ANNOTATIONS", "$CONSTRUCTOR_REFERENCES", "$PROPERTY_REFERENCES", "$METHOD_REFERENCES")
                .build());

        // 重载load方法
        builder.addMethod(MethodSpec.methodBuilder("load")
                .addAnnotation(ctx.getClassName(Override.class))
                .addModifiers(Modifier.PUBLIC)
                .returns(ParameterizedTypeName.get(ctx.getClassName(BeanMirror.class), ctx.getClassName(declaredType)))
                .addStatement("return this")
                .build());

        // dispatchInstance方法
        BuildFormat dispatchInstance = resolveDispatchInstance(declaredType, constructors, ctx);
        builder.addMethod(MethodSpec.methodBuilder("dispatchInstance")
                .addAnnotation(ctx.getClassName(Override.class))
                .addModifiers(Modifier.PROTECTED)
                .returns(ClassName.get(declaredType))
                .addParameter(ParameterSpec.builder(TypeName.INT, "index").build())
                .addParameter(ParameterSpec.builder(ArrayTypeName.of(TypeName.OBJECT), "args").build())
                .addCode(dispatchInstance.format(), dispatchInstance.args())
                .build());

        // dispatchMethod
        BuildFormat dispatchMethod = resolveDispatchMethod(declaredType, methods.values(), ctx);
        builder.addMethod(MethodSpec.methodBuilder("dispatchMethod")
                .addAnnotation(ctx.getClassName(Override.class))
                .addAnnotation(AnnotationSpec.builder(SuppressWarnings.class)
                        .addMember("value", "$S", "unchecked")
                        .build())
                .addModifiers(Modifier.PROTECTED)
                .returns(TypeName.OBJECT)
                .addParameter(ParameterSpec.builder(TypeName.INT, "index").build())
                .addParameter(ParameterSpec.builder(TypeName.OBJECT, "target").build())
                .addParameter(ParameterSpec.builder(ArrayTypeName.of(TypeName.OBJECT), "args").build())
                .addCode(dispatchMethod.format(), dispatchMethod.args())
                .build());

        TypeSpec typeSpec = builder.build();

        // 生成类文件
        ctx.writeJavaFile(ctx.getPackageNameOf(declaredType), typeSpec);

        // 生成配置文件
        ctx.writeMetaInf(BeanReference.class, "%s$Ref".formatted(ctx.getFullyQualifiedName(declaredType)));
    }

    private List<Element> resolveEnclosedElements(DeclaredType declaredType) {
        TypeElement element = (TypeElement) declaredType.asElement();
        List<Element> enclosedElements = new ArrayList<>(element.getEnclosedElements());
        TypeMirror superclass = element.getSuperclass();
        while (superclass != null && !superclass.toString().equals("java.lang.Object")) {
            TypeElement superElement = (TypeElement) ((DeclaredType) superclass).asElement();
            enclosedElements.addAll(superElement.getEnclosedElements());
            superclass = superElement.getSuperclass();
        }
        return enclosedElements;
    }

    private List<AnnotationMirror> resolveAnnotationMirrors(DeclaredType declaredType, ProcessorContext ctx) {
        List<AnnotationMirror> annotationMirrors = new ArrayList<>();
        List<? extends AnnotationMirror> declared = declaredType.asElement().getAnnotationMirrors();
        for (AnnotationMirror annotationMirror : declared) {
            if (ctx.isMirrorAnnotation(annotationMirror)) {
                annotationMirrors.add(annotationMirror);
            }
        }
        TypeMirror superclass = ((TypeElement) declaredType.asElement()).getSuperclass();
        while (superclass != null && !superclass.toString().equals("java.lang.Object")) {
            List<? extends AnnotationMirror> annotations = ((DeclaredType) superclass).asElement().getAnnotationMirrors();
            for (AnnotationMirror annotation : annotations) {
                if (ctx.isInheritedAnnotated(annotation)
                        && ctx.isMirrorAnnotation(annotation)) {
                    annotationMirrors.add(annotation);
                }
            }
            superclass = ((TypeElement) ((DeclaredType) superclass).asElement()).getSuperclass();
        }
        return annotationMirrors;
    }

    private BuildFormat resolveDeclaredAnnotations(String tabs, List<AnnotationMirror> annotationMirrors, ProcessorContext ctx) {
        StringBuilder format = new StringBuilder("new $T[] {");
        if (!annotationMirrors.isEmpty()) {
            format.append("\n");
        }
        List<Object> args = new ArrayList<>();
        args.add(ctx.getClassName(jamirr.AnnotationMirror.class));
        boolean first1 = true;
        for (AnnotationMirror annotation : annotationMirrors) {
            if (!first1) {
                format.append(",\n");
            }
            format.append(tabs);
            format.append("\t$T.load($S).build(");
            args.add(ctx.getClassName(AnnotationMirrorBuilder.class));
            args.add(annotation.getAnnotationType().toString());
            boolean first = true;
            for (Map.Entry<? extends ExecutableElement, ? extends AnnotationValue> pair : annotation.getElementValues().entrySet()) {
                if (!first) {
                    format.append(", ");
                }
                format.append("$S");
                args.add(pair.getKey().getSimpleName().toString());
                format.append(", $L");
                args.add(pair.getValue());
                first = false;
            }
            format.append(")");
            first1 = false;
        }
        if (!annotationMirrors.isEmpty()) {
            format.append("\n");
            format.append(tabs);
        }
        format.append("}");

        return new BuildFormat(format.toString(), args.toArray());
    }

    private BuildFormat resolveDeclaredAnnotations(DeclaredType declaredType, ProcessorContext ctx) {
        List<AnnotationMirror> annotationMirrors = resolveAnnotationMirrors(declaredType, ctx);
        return resolveDeclaredAnnotations("", annotationMirrors, ctx);
    }

    private BuildFormat resolveDispatchInstance(DeclaredType declaredType, Collection<DeclaredConstructor> constructors, ProcessorContext ctx) {
        StringBuilder sb = new StringBuilder("return switch (index) {\n");
        List<Object> args = new ArrayList<>();
        for (DeclaredConstructor ctor : constructors) {
            sb.append("\tcase $L -> new $T(");
            args.add(ctor.index());
            args.add(ctx.getClassName(declaredType));
            int count = 0;
            boolean first = true;
            for (VariableElement parameter : ctor.element().getParameters()) {
                if (!first) {
                    sb.append(", ");
                }
                sb.append("($T) args[$L]");
                args.add(parameter.asType());
                args.add(count++);
                first = false;
            }
            sb.append(");\n");
        }
        sb.append("\tdefault -> throw new IllegalArgumentException(\"Invalid constructor index: \" + index);\n");
        sb.append("};");
        return new BuildFormat(sb.toString(), args.toArray());
    }

    private BuildFormat resolveDispatchMethod(DeclaredType declaredType, Collection<DeclaredMethod> methods, ProcessorContext ctx) {
        StringBuilder sb = new StringBuilder("return switch (index) {\n");
        List<Object> args = new ArrayList<>();
        for (DeclaredMethod method : methods) {
            boolean isVoid = "void".equals(method.returnType());
            sb.append("\tcase $L -> ");
            if (isVoid) {
                sb.append("{\n\t\t");
            }
            sb.append("(($T) target).$L(");
            args.add(method.index());
            args.add(ctx.getClassName(declaredType));
            args.add(method.element().getSimpleName().toString());
            int count = 0;
            boolean first = true;
            for (VariableElement parameter : method.element().getParameters()) {
                if (!first) {
                    sb.append(", ");
                }
                sb.append("($T) args[$L]");
                args.add(parameter.asType());
                args.add(count++);
                first = false;
            }
            sb.append(");\n");
            if (isVoid) {
                sb.append("\t\tyield null;\n");
                sb.append("\t}\n");
            }
        }
        sb.append("\tdefault -> throw new IllegalArgumentException(\"Invalid method index: \" + index);\n");
        sb.append("};");
        return new BuildFormat(sb.toString(), args.toArray());
    }

    @SuppressWarnings("unchecked")
    private BuildFormat resolvePropertyReferences(Collection<Element> properties, Map<String, DeclaredMethod> methods, ProcessorContext ctx) {
        StringBuilder sb = new StringBuilder("new PropertyRef[] {\n");
        List<Object> args = new ArrayList<>();
        boolean first = true;
        for (Element prop : properties) {
            if (!first) {
                sb.append(",\n");
            }
            sb.append("\tnew PropertyRef(\n\t\tArgument.of($T.class, $S)");
            args.add(prop.asType());
            String propName = prop.getSimpleName().toString();
            args.add(propName);

            int readMethodIndex = findReadMethodIndex(propName, prop.asType().toString(), methods);
            int writeMethodIndex = findWriteMethodIndex(propName, methods);

            sb.append(",\n\t\t$L");// readable
            args.add(readMethodIndex >= 0);
            sb.append(",\n\t\t$L");// getMethodIndex
            args.add(readMethodIndex);

            sb.append(",\n\t\t$L");// writable
            args.add(writeMethodIndex >= 0);
            sb.append(",\n\t\t$L");// setMethodIndex
            args.add(writeMethodIndex);

            BuildFormat format = resolveDeclaredAnnotations("\t\t", (List<AnnotationMirror>) prop.getAnnotationMirrors(), ctx);
            sb.append(",\n\t\t");
            sb.append(format.format());
            args.addAll(Arrays.asList(format.args()));
            sb.append(")");
            first = false;
        }
        sb.append("\n}");
        return new BuildFormat(sb.toString(), args.toArray());
    }

    @SuppressWarnings("unchecked")
    private BuildFormat resolveConstructorReferences(Collection<DeclaredConstructor> constructors, ProcessorContext ctx) {
        StringBuilder sb = new StringBuilder("new ConstructorRef[] {\n");
        List<Object> args = new ArrayList<>();
        boolean first = true;
        for (DeclaredConstructor ctor : constructors) {
            if (!first) {
                sb.append(",\n");
            }
            sb.append("\tnew ConstructorRef(\n\t\t$L");
            args.add(ctor.index());
            if (ctor.element().getParameters().isEmpty()) {
                sb.append(",\n\t\tnull");
            } else {
                sb.append(",\n\t\t");
                resolveArguments(sb, args, ctor.element(), ctx);
            }
            BuildFormat format = resolveDeclaredAnnotations("\t\t", (List<AnnotationMirror>) ctor.element().getAnnotationMirrors(), ctx);
            sb.append(",\n\t\t");
            sb.append(format.format());
            args.addAll(Arrays.asList(format.args()));
            sb.append(")");
            first = false;
        }
        sb.append("\n}");
        return new BuildFormat(sb.toString(), args.toArray());
    }

    @SuppressWarnings("unchecked")
    private BuildFormat resolveMethodReferences(Collection<DeclaredMethod> methods, ProcessorContext ctx) {
        StringBuilder sb = new StringBuilder("new MethodRef[] {\n");
        List<Object> args = new ArrayList<>();
        boolean first = true;
        for (DeclaredMethod method : methods) {
            if (!first) {
                sb.append(",\n");
            }
            sb.append("\tnew MethodRef(\n\t\t$L");
            args.add(method.index());
            if ("void".equals(method.returnType())) {
                sb.append(",\n\t\tnull");
            } else {
                sb.append(",\n\t\tArgument.of($T.class)");
                args.add(method.element().getReturnType());
            }
            sb.append(",\n\t\t$S");
            args.add(method.element().getSimpleName().toString());
            if (method.element().getParameters().isEmpty()) {
                sb.append(",\n\t\tnull");
            } else {
                sb.append(",\n\t\t");
                resolveArguments(sb, args, method.element(), ctx);
            }
            BuildFormat format = resolveDeclaredAnnotations("\t\t", (List<AnnotationMirror>) method.element().getAnnotationMirrors(), ctx);
            sb.append(",\n\t\t");
            sb.append(format.format());
            args.addAll(Arrays.asList(format.args()));
            sb.append(")");
            first = false;
        }
        sb.append("\n}");
        return new BuildFormat(sb.toString(), args.toArray());
    }

    private void resolveArguments(StringBuilder sb, List<Object> args, ExecutableElement method, ProcessorContext ctx) {
        sb.append("new $T[]{");
        args.add(ctx.getClassName(Argument.class));
        boolean first = true;
        for (VariableElement param : method.getParameters()) {
            if (!first) {
                sb.append(", ");
            }
            sb.append("Argument.of($T.class)");
            args.add(param.asType());
            first = false;
        }
        sb.append("}");
    }

    private int findReadMethodIndex(String name, String returnType, Map<String, DeclaredMethod> methods) {
        String methodName = ("boolean".equals(returnType) ? "is" : "get") + capitalize(name);
        DeclaredMethod method = methods.get(methodName);
        if (method == null) {
            method = methods.get(name);
            if (method == null) {
                return -1;
            }
        }
        if (method.args() > 0) {
            return -1;
        }
        return method.index();
    }

    private int findWriteMethodIndex(String name, Map<String, DeclaredMethod> methods) {
        String methodName = "set" + capitalize(name);
        DeclaredMethod method = methods.get(methodName);
        if (method == null) {
            method = methods.get(name);
            if (method == null) {
                return -1;
            }
        }
        if (method.args() != 1) {
            return -1;
        }
        return method.index();
    }

    public static String capitalize(String str) {
        if (str == null || str.isEmpty()) {
            return str;
        }

        char baseChar = str.charAt(0);
        char updatedChar = Character.toUpperCase(baseChar);
        if (baseChar == updatedChar) {
            return str;
        }

        char[] chars = str.toCharArray();
        chars[0] = updatedChar;
        return new String(chars);
    }

    private record BuildFormat(String format, Object[] args) {
    }

    private record DeclaredConstructor(int index, ExecutableElement element) {
    }

    private record DeclaredMethod(int index, String returnType, int args, String argType, ExecutableElement element) {
    }
}
